import json
from abc import abstractmethod
from functools import cached_property
from pathlib import Path
from typing import Any

import numpy as np
from numpy.typing import NDArray
from PIL import Image

from immich_ml.config import log
from immich_ml.models.base import InferenceModel
from immich_ml.models.transforms import (
    crop_pil,
    decode_pil,
    get_pil_resampling,
    normalize,
    resize_pil,
    serialize_np_array,
    to_numpy,
)
from immich_ml.schemas import ModelSession, ModelTask, ModelType


class BaseCLIPVisualEncoder(InferenceModel):
    """
    CLIP视觉编码器的基础抽象类
    用于将图像转换为向量表示，主要用于图像搜索任务
    """
    depends = []
    identity = (ModelType.VISUAL, ModelTask.SEARCH)

    def _predict(self, inputs: Image.Image | bytes, **kwargs: Any) -> str:
        """
        预测方法：将输入图像转换为向量表示
        Args:
            inputs: 输入图像（PIL Image或字节数据）
        Returns:
            序列化后的numpy数组（向量表示）
        """
        log.info(f"开始处理图像预测，模型名称: {self.model_name}")
        image = decode_pil(inputs)
        log.debug(f"图像解码完成，尺寸: {image.size}")
        res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0]
        log.debug(f"模型推理完成，输出向量维度: {res.shape}")
        return serialize_np_array(res)

    @abstractmethod
    def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
        """
        抽象方法：定义图像转换的具体实现
        Args:
            image: 输入图像
        Returns:
            转换后的图像数据字典
        """
        pass

    @property
    def model_cfg_path(self) -> Path:
        """模型配置文件路径"""
        path = self.cache_dir / "config.json"
        log.debug(f"模型配置文件路径: {path}")
        return path

    @property
    def preprocess_cfg_path(self) -> Path:
        """预处理配置文件路径"""
        path = self.model_dir / "preprocess_cfg.json"
        log.debug(f"预处理配置文件路径: {path}")
        return path

    @cached_property
    def model_cfg(self) -> dict[str, Any]:
        """加载模型配置"""
        log.info(f"开始加载CLIP模型配置，模型名称: '{self.model_name}'")
        model_cfg: dict[str, Any] = json.load(self.model_cfg_path.open())
        log.info(f"CLIP模型配置加载完成，模型名称: '{self.model_name}'")
        log.debug(f"模型配置内容: {model_cfg}")
        return model_cfg

    @cached_property
    def preprocess_cfg(self) -> dict[str, Any]:
        """加载预处理配置"""
        log.info(f"开始加载CLIP模型预处理配置，模型名称: '{self.model_name}'")
        preprocess_cfg: dict[str, Any] = json.load(self.preprocess_cfg_path.open())
        log.info(f"CLIP模型预处理配置加载完成，模型名称: '{self.model_name}'")
        log.debug(f"预处理配置内容: {preprocess_cfg}")
        return preprocess_cfg


class OpenClipVisualEncoder(BaseCLIPVisualEncoder):
    """
    OpenCLIP视觉编码器的具体实现
    实现了图像预处理和转换的具体逻辑
    """
    def _load(self) -> ModelSession:
        """
        加载模型并初始化预处理参数
        Returns:
            模型会话对象
        """
        log.info(f"开始加载OpenCLIP视觉编码器，模型名称: {self.model_name}")
        size: list[int] | int = self.preprocess_cfg["size"]
        self.size = size[0] if isinstance(size, list) else size
        log.debug(f"设置图像处理尺寸: {self.size}")

        self.resampling = get_pil_resampling(self.preprocess_cfg["interpolation"])
        self.mean = np.array(self.preprocess_cfg["mean"], dtype=np.float32)
        self.std = np.array(self.preprocess_cfg["std"], dtype=np.float32)
        log.debug(f"预处理参数 - 重采样方法: {self.resampling}, 均值: {self.mean}, 标准差: {self.std}")

        session = super()._load()
        log.info(f"OpenCLIP视觉编码器加载完成，模型名称: {self.model_name}")
        return session

    def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
        """
        图像转换方法：调整大小、裁剪、标准化
        Args:
            image: 输入图像
        Returns:
            处理后的图像数据字典
        """
        log.debug(f"开始图像转换，原始尺寸: {image.size}")
        image = resize_pil(image, self.size)
        log.debug(f"图像调整大小完成，新尺寸: {image.size}")
        
        image = crop_pil(image, self.size)
        log.debug(f"图像裁剪完成，尺寸: {image.size}")
        
        image_np = to_numpy(image)
        log.debug(f"转换为numpy数组，形状: {image_np.shape}")
        
        image_np = normalize(image_np, self.mean, self.std)
        log.debug(f"图像标准化完成，数据范围: [{image_np.min()}, {image_np.max()}]")
        
        result = {"image": np.expand_dims(image_np.transpose(2, 0, 1), 0)}
        log.debug(f"最终输出张量形状: {result['image'].shape}")
        return result
